#ifndef CMD_CC
#define CMD_CC

#include<string>
#include<fstream>
#include<sstream>
#include<iostream>
#include<algorithm>
#include<iterator>
#include<vector>
#include<map>
#include<cstdlib>
#include<ctime>
#include<cmath>
#include<tclap/CmdLine.h>
#include "def.hpp"
#include "cmd.hpp"

using namespace TCLAP; // using the command line parser                                          
using namespace std;

PuzaCMD* PuzaCMD::m_pInstance = NULL;

PuzaCMD* PuzaCMD::Instance(){
  if (!m_pInstance)   // Only allow one instance of class to be generated.                                                                             
    m_pInstance = new PuzaCMD;
  return m_pInstance;
}

PuzaCMD::PuzaCMD(){ }


ModelType PuzaCMD::checkModel(string temp_model){
  if(temp_model.compare("topicmodel")==0)
    return MODEL_TOPICMODEL;
  if(temp_model.compare("generator")==0)
    return MODEL_GENERATOR;
  if(temp_model.compare("estimator")==0)
    return MODEL_ESTIMATOR;
  return MODEL_ERROR;
}

GeneratorType PuzaCMD::checkGenerator(string temp_generator){
  if(temp_generator.compare("multinomial")==0)
    return GE_MULTINOMIAL;
  if(temp_generator.compare("dirichlet")==0)
    return GE_DIRICHLET;
  if(temp_generator.compare("polya")==0)
    return GE_POLYA;
  if(temp_generator.compare("error")==0)
    return GE_ERROR;
  return GE_ERROR;
}

EstimatorType PuzaCMD::checkEstimator(string temp_estimator){
  if(temp_estimator.compare("polya")==0)
    return EST_POLYA;
  return EST_ERROR;
}

PolyaEstimatorType PuzaCMD::checkPolyaEstimator(string temp_polya_estimator){
  if(temp_polya_estimator.compare("sgd")==0)
    return EST_POLYA_SGD;
  if(temp_polya_estimator.compare("newton")==0)
    return EST_POLYA_NEWTON;
  if(temp_polya_estimator.compare("fixedpoint")==0)
    return EST_POLYA_FIXEDPOINT;
  if(temp_polya_estimator.compare("error")==0)
    return EST_POLYA_ERROR;
  return EST_POLYA_ERROR;
}

TopicModelType PuzaCMD::checkTopicModelType(string temp_topicmodel_type){
  if(temp_topicmodel_type.compare("lda_gibbs")==0)
    return TOPICMODEL_LDA_GIBBS_SAMPLING;
  if(temp_topicmodel_type.compare("lda_vb")==0)
    return TOPICMODEL_LDA_VARIATIONAL_BAYES;
  if(temp_topicmodel_type.compare("lda_vbonline")==0)
    return TOPICMODEL_LDA_VARIATIONAL_BAYES_ONLINE;
  if(temp_topicmodel_type.compare("tot")==0)
    return TOPICMODEL_TOT;
  if(temp_topicmodel_type.compare("geo_simple_1_batch")==0)
    return TOPICMODEL_GEO_SIMPLE_1_BATCH;
  if(temp_topicmodel_type.compare("geo_simple_1_online")==0)
    return TOPICMODEL_GEO_SIMPLE_1_ONLINE;
  if(temp_topicmodel_type.compare("geo_cluster_batch")==0)
    return TOPICMODEL_GEO_CLUSTER_BATCH;
  if(temp_topicmodel_type.compare("geo_mf_batch")==0)
    return TOPICMODEL_GEO_MF_BATCH;
  if(temp_topicmodel_type.compare("geo_mf_topic")==0)
    return TOPICMODEL_GEO_MF_TOPIC;
  if(temp_topicmodel_type.compare("geo_mf_topic_back")==0)
    return TOPICMODEL_GEO_MF_TOPIC_BACK;
  if(temp_topicmodel_type.compare("geo_mf_topic_geo_back")==0)
    return TOPICMODEL_GEO_MF_TOPIC_GEO_BACK;
  if(temp_topicmodel_type.compare("geo_mf_online")==0)
    return TOPICMODEL_GEO_MF_ONLINE;
  if(temp_topicmodel_type.compare("geo_mf_loc")==0)
    return TOPICMODEL_GEO_MF_LOC;
  if(temp_topicmodel_type.compare("geo_mf_loc_back")==0)
    return TOPICMODEL_GEO_MF_LOC_BACK;
  if(temp_topicmodel_type.compare("at_batch")==0)
    return TOPICMODEL_AT_BATCH;
  if(temp_topicmodel_type.compare("at_online")==0)
    return TOPICMODEL_AT_ONLINE;
  return TOPICMODEL_ERROR;
}

TopicModelTaskType PuzaCMD::checkTopicModelTaskType(string temp_task){
  if(temp_task.compare("infer")==0)
    return TOPICMODEL_TASK_INFERENCE;
  if(temp_task.compare("predict")==0)
    return TOPICMODEL_TASK_PREDICTION;
  if(temp_task.compare("infer_predict")==0)
    return TOPICMODEL_TASK_INFERENCE_AND_PREDICTION;
  if(temp_task.compare("interactive")==0)
    return TOPICMODEL_TASK_INTERACTIVE;
  return TOPICMODEL_TASK_ERROR;
}

void PuzaCMD::parseCMD(int argc, char* argv[]){

  CmdLine cmd("This is the program implementing multiple text mining algorithms, including topic models.", ' ', "0.5");

  /* Main Program Parameters */
  ValueArg<string> modelArg("","model","model name",true,"","string");
  cmd.add( modelArg );

  /* IO Parameters */
  ValueArg<string> inputArg("","input","input file",false,"","string");
  cmd.add( inputArg );
  ValueArg<string> trainArg("","train","the training file",false,"","string");
  cmd.add( trainArg );
  ValueArg<string> testArg("","test","the testing file",false,"","string");
  cmd.add( testArg );  
  ValueArg<string> outputArg("","output","output file",false,"","string");
  cmd.add( outputArg );
  ValueArg<string> input_format_Arg("","inputformat","input format",false,"text","string");
  ValueArg<string> output_format_Arg("","outputformat","output format",false,"text","string");

  /* Topic Model Parameters */
  ValueArg<string> topicModelArg("","topicmodel", "topic model name",false,"error","string");
  cmd.add(topicModelArg);
  ValueArg<string> topicModelTaskArg("","task", "the topic model task",false,"error","string");
  cmd.add(topicModelTaskArg);
  ValueArg<int> DocNUMArg("","docs", "the number of documents",false,-1,"int");
  cmd.add(DocNUMArg);
  ValueArg<int> TermNUMArg("","terms", "the number of distinct terms",false,-1,"int");
  cmd.add(TermNUMArg);
  ValueArg<int> AuthorNUMArg("","authors", "the number of distinct authors",false,-1,"int");
  cmd.add(AuthorNUMArg);
  ValueArg<int> TopicNUMArg("","topics", "the number of topics",false,-1,"int");
  cmd.add(TopicNUMArg);
  ValueArg<int> BatchSizeArg("","batch", "the size of mini-batch",false,100,"int");
  cmd.add(BatchSizeArg);
  ValueArg<double> LDAInitialAlphaArg("","lda_alpha", "initial alpha values",false,0.5,"double");
  cmd.add(LDAInitialAlphaArg);
  ValueArg<double> LDAInitialBetaArg("","lda_beta", "initial beta values",false,0.01,"double");
  cmd.add(LDAInitialBetaArg);
  ValueArg<double> LDAVBONLINE_tau0_Arg("","tau0", "tau0 for VB Online",false,1024.0,"double");
  cmd.add(LDAVBONLINE_tau0_Arg);
  ValueArg<double> LDAVBONLINE_kappa_Arg("","kappa", "kappa for VB online",false,0.5,"double");
  cmd.add(LDAVBONLINE_kappa_Arg);
  ValueArg<int> GeoValueTypeArg("","geo_type","the type of geo value",false,0,"int");
  cmd.add(GeoValueTypeArg);
  ValueArg<int> GeoRegionNumArg("","geo_region_num","the number of geo regions",false,0,"int");
  cmd.add(GeoRegionNumArg);
  ValueArg<int> GeoUserNumArg("","geo_user_num","the number of users",false,0,"int");
  cmd.add(GeoUserNumArg);
  ValueArg<double> GeoLearningRateArg("","geo_learning_rate","the learning rate in the m-step",false,0.0001,"double");
  cmd.add(GeoLearningRateArg);
  ValueArg<double> GeoRegArg("","geo_reg","the regularization parameter",false,1.0,"double");
  cmd.add(GeoRegArg);
  ValueArg<int> GeoRegL1KArg("","geo_reg_l1_k","the regularization parameter K for L1",false,20,"int");
  cmd.add(GeoRegL1KArg);
  ValueArg<double> GeoRegL1GArg("","geo_reg_l1_g","the regularization parameter g for L1",false,0.1,"double");
  cmd.add(GeoRegL1GArg);
  ValueArg<string> GeoRegTypeArg("","geo_reg_type","the regularization type",false,"l2","string");
  cmd.add(GeoRegTypeArg);
  ValueArg<int> GeoMIterArg("","geo_m_iter","the number of times M step is performed",false,1,"int");
  cmd.add(GeoMIterArg);
  ValueArg<int> GeoEIterArg("","geo_e_iter","the number of times E step is performed",false,1,"int");
  cmd.add(GeoEIterArg);

  /* Estimator Parameters */
  ValueArg<string> estimatorArg("","estimator","estimator name",false,"error","string");
  cmd.add( estimatorArg );
  ValueArg<string> polyaEstimatorArg("","polyaEstimator","estimator for Polya distribution (Dirichlet/Multinomial)",false,"error","string");
  cmd.add( polyaEstimatorArg );
  ValueArg<int> estimator_sample_Arg("","estimatorSamples","the number of samples to be estimated",false,-1,"int");
  cmd.add(estimator_sample_Arg);
  ValueArg<double> polyaSGDRateArg("","polyaSGDRate","the initial learning rate of Stochastic Gradient Descent for Polya distribution.",false,1e-5,"double");
  cmd.add(polyaSGDRateArg);

  /* Generator Parameters */
  ValueArg<string> generatorArg("","generator","generator name",false,"error","string");
  cmd.add( generatorArg );

  ValueArg<int> multinomial_k_Arg("","multinomialK","the dimension of a multinomial distribution",false,-1,"int");
  cmd.add(multinomial_k_Arg);
  ValueArg<int> multinomial_n_Arg("","multinomialN","the number of trials of a multinomial distribution",false,-1,"int");
  cmd.add(multinomial_n_Arg);
  MultiArg<double> multinomial_p_Arg("","multinomialP","each element of the multinomial distribution", false,"double" );
  cmd.add(multinomial_p_Arg);

  ValueArg<int> dirichlet_k_Arg("","dirichletK","the dimension of a dirichlet distribution",false,-1,"int");
  cmd.add(dirichlet_k_Arg);  
  ValueArg<double> dirichlet_alpha_Arg("","dirichletAlpha","the scalar of a dirichlet distribution",false,1.0,"double");
  cmd.add(dirichlet_alpha_Arg);  
  MultiArg<double> dirichlet_base_Arg("","dirichletBase","each element of the base measure of dirichlet distribution", false,"double" );
  cmd.add(dirichlet_base_Arg);

  ValueArg<int> polya_k_Arg("","polyaK","the dimension of a polya distribution",false,-1,"int");
  cmd.add(polya_k_Arg);
  ValueArg<int> polya_n_Arg("","polyaN","the number of trials of a polya distribution",false,-1,"int");
  cmd.add(polya_n_Arg);
  ValueArg<double> polya_alpha_Arg("","polyaAlpha","the scalar of a polya distribution",false,1.0,"double");
  cmd.add(polya_alpha_Arg);
  MultiArg<double> polya_base_Arg("","polyaBase","each element of the base measure of polya distribution", false,"double" );
  cmd.add(polya_base_Arg);  

  ValueArg<int> generator_sample_Arg("","generatorSamples","the number of samples drawn from a generator",false,1,"int");
  cmd.add(generator_sample_Arg); 

  /* Parse */
  cmd.parse( argc, argv );

  /* Get Main Program Parameters */
  string temp_model = modelArg.getValue();
  PuzaDEF::Instance()->currentModelType = checkModel(temp_model); // check whether we support the model
  if(PuzaDEF::Instance()->currentModelType == MODEL_ERROR){
    cout << "We currently only support 1) Synthetic Data Model(generator) 2) Topic Models(topicmodel) 3) Estimators(estimator)." << endl;
    exit(0);
  }

  PuzaDEF::Instance()->input_file_name = inputArg.getValue();
  PuzaDEF::Instance()->train_file_name = trainArg.getValue();
  PuzaDEF::Instance()->test_file_name = testArg.getValue();
  PuzaDEF::Instance()->output_file_name = outputArg.getValue();

  /* Get Topic Model Parameters */
  if(PuzaDEF::Instance()->currentModelType == MODEL_TOPICMODEL){
    string temp_topicmodel = topicModelArg.getValue();
    PuzaDEF::Instance()->currentTopicModelType = checkTopicModelType(temp_topicmodel);
    if(PuzaDEF::Instance()->currentTopicModelType == TOPICMODEL_ERROR){
      cout << "We currently only support: \n \t1) Latent Dirichlet Allocation - Gibbs Sampling (lda_gibbs). \
\n\t2) Latent Dirichlet Allocation - Variatinoal Bayes (lda_vb). \
\n\t3) Latent Dirichlet Allocation - Online Variatinoal Bayes (lda_vbonline). \
\n\t4) Topics over Time (tot). \
\n\t5) Geo Topic Model - Simple Model I (geo_simple_1). \
" << endl;
      exit(0);
    }
    // get the task to be performed
    string temp_task = topicModelTaskArg.getValue();
    PuzaDEF::Instance()->currentTopicModelTaskType = checkTopicModelTaskType(temp_task);
    if(PuzaDEF::Instance()->currentTopicModelTaskType == TOPICMODEL_TASK_ERROR){
      cout << "We currently only support 1) Inference (infer) 2) Prediction (predict) 3) Perplexity (perplexity). " << endl;
      exit(0);
    }
    if(PuzaDEF::Instance()->currentTopicModelTaskType == TOPICMODEL_TASK_INFERENCE){
      if(PuzaDEF::Instance()->input_file_name.empty()){
	cout << "Please specify the input file by using the option [input]" << endl;
	exit(0);
      }
      if(PuzaDEF::Instance()->output_file_name.empty()){
	cout << "Please specify the output file by using the option [output]" << endl;
	exit(0);
      }
    }
    PuzaDEF::Instance()->DOC_NUM = DocNUMArg.getValue();
    PuzaDEF::Instance()->TERM_NUM = TermNUMArg.getValue();
    PuzaDEF::Instance()->TOPIC_NUM = TopicNUMArg.getValue();
    PuzaDEF::Instance()->AUTHOR_NUM = AuthorNUMArg.getValue();
    PuzaDEF::Instance()->LDA_initial_alpha = LDAInitialAlphaArg.getValue();
    PuzaDEF::Instance()->LDA_initial_beta = LDAInitialBetaArg.getValue();      
    PuzaDEF::Instance()->LDA_VBONLINE_tau0 = LDAVBONLINE_tau0_Arg.getValue();
    PuzaDEF::Instance()->LDA_VBONLINE_kappa = LDAVBONLINE_kappa_Arg.getValue();
    PuzaDEF::Instance()->BatchSize = BatchSizeArg.getValue();
    if(GeoValueTypeArg.getValue() == 1)
      PuzaDEF::Instance()->GEO_TYPE = GEO_TYPE_DISCRETE;
    PuzaDEF::Instance()->GEO_USER_NUM = GeoUserNumArg.getValue();
    PuzaDEF::Instance()->GEO_REGION_NUM = GeoRegionNumArg.getValue();
    PuzaDEF::Instance()->GEO_LEARNING_RATE = GeoLearningRateArg.getValue();
    PuzaDEF::Instance()->GEO_REG = GeoRegArg.getValue();
    PuzaDEF::Instance()->GEO_M_ITER = GeoMIterArg.getValue();
    PuzaDEF::Instance()->GEO_E_ITER = GeoEIterArg.getValue();
    string tempRegType = GeoRegTypeArg.getValue();
    if(tempRegType.compare("l2")==0){
      PuzaDEF::Instance()->GEO_REG_TYPE = 0; // L2 regularization
    }
    else{
      PuzaDEF::Instance()->GEO_REG_TYPE = 1; // L1 regularization
    }
    PuzaDEF::Instance()->GEO_REG_L1_K = GeoRegL1KArg.getValue();
    PuzaDEF::Instance()->GEO_REG_L1_G = GeoRegL1GArg.getValue();
    
  }

  /* Get Estimator Parameters */
  if(PuzaDEF::Instance()->currentModelType == MODEL_ESTIMATOR){
    string temp_estimator = estimatorArg.getValue();
    PuzaDEF::Instance()->currentEstimatorType = checkEstimator(temp_estimator);
    PuzaDEF::Instance()->EST_Samples = estimator_sample_Arg.getValue();
    
    if(PuzaDEF::Instance()->currentEstimatorType == EST_ERROR){
      cout << "We currently only support to estimate 1) Polya distribution (polya). " << endl;
      exit(0);
    }
    
    if(PuzaDEF::Instance()->currentEstimatorType == EST_POLYA){
      string temp_polya_estimator = polyaEstimatorArg.getValue();
      PuzaDEF::Instance()->currentPolyaEstimatorType = checkPolyaEstimator(temp_polya_estimator);
      PuzaDEF::Instance()->EST_Polya_K = polya_k_Arg.getValue();    
      PuzaDEF::Instance()->EST_Polya_SGD_Initial_Rate = polyaSGDRateArg.getValue();
      if ((PuzaDEF::Instance()->EST_Polya_K == -1) || (PuzaDEF::Instance()->EST_Samples == -1)){
	cout << "In order to estimate polya distribution, we need K and sample numbers." << endl;
	exit(0);
      }
      if(PuzaDEF::Instance()->currentPolyaEstimatorType == EST_POLYA_ERROR){
	cout << "We currently only support to use 1) Stochastic Gradient Descent (sgd) 2) Newton (newton) 3) Fixed Point Iteration (fixedpoint) estimate Polya distirbution. " << endl;
	exit(0);
      }
    }
  }


  /* Get Generator Parameters */

  if(PuzaDEF::Instance()->currentModelType == MODEL_GENERATOR){
    // get the type of generator
    string temp_generator = generatorArg.getValue();
    PuzaDEF::Instance()->currentGeneratorType = checkGenerator(temp_generator); // check whether we support the generator
    if(PuzaDEF::Instance()->currentGeneratorType == GE_ERROR){
      cout << "We currently only support 1) Multinomial (multinomial) 2) Dirichlet (dirichlet) 3) Polya -- Dirichlet/Multinomial (polya)." << endl;
      exit(0);
    }


    PuzaDEF::Instance()->GE_Samples = generator_sample_Arg.getValue();

    if(PuzaDEF::Instance()->currentGeneratorType == GE_POLYA){
      PuzaDEF::Instance()->GE_Polya_K = polya_k_Arg.getValue();
      PuzaDEF::Instance()->GE_Polya_N = polya_n_Arg.getValue();
      PuzaDEF::Instance()->GE_Polya_Alpha = polya_alpha_Arg.getValue();

      vector<double> raw_polya_base = polya_base_Arg.getValue();
      if ((PuzaDEF::Instance()->GE_Polya_K == -1) || (PuzaDEF::Instance()->GE_Polya_N == -1)){
	cout << "In order to generate polya random variables, we need K and N. (Assume scalar is 1.0)" << endl;
	exit(0);
      }

      if(raw_polya_base.size()!=0){
	PuzaDEF::Instance()->GE_Polya_Base = new double[raw_polya_base.size()]; // allocate memeory
	for ( int i = 0; static_cast<unsigned int>(i) < raw_polya_base.size(); i++ ){
	  PuzaDEF::Instance()->GE_Polya_Base[i] = raw_polya_base[i];
	}
      }
      else{
	PuzaDEF::Instance()->GE_Polya_Base = new double[PuzaDEF::Instance()->GE_Polya_K];
	for (int k=0;k < PuzaDEF::Instance()->GE_Polya_K; k++){
	  PuzaDEF::Instance()->GE_Polya_Base[k] = 1.0/PuzaDEF::Instance()->GE_Polya_K;
	}
      }
    }
    if(PuzaDEF::Instance()->currentGeneratorType == GE_MULTINOMIAL){
      PuzaDEF::Instance()->GE_Multinomial_K = multinomial_k_Arg.getValue();
      PuzaDEF::Instance()->GE_Multinomial_N = multinomial_n_Arg.getValue();

      vector<double> raw_multinomial_p = multinomial_p_Arg.getValue();
      if ((PuzaDEF::Instance()->GE_Multinomial_K == -1) || (PuzaDEF::Instance()->GE_Multinomial_N == -1) || (raw_multinomial_p.size() == 0)){
	cout << "In order to generate multinomial random variables, we need K, N and P." << endl;
	exit(0);
      }

      PuzaDEF::Instance()->GE_Multinomial_P = new double[raw_multinomial_p.size()]; // allocate memeory
      for ( int i = 0; static_cast<unsigned int>(i) < raw_multinomial_p.size(); i++ ){
	PuzaDEF::Instance()->GE_Multinomial_P[i] = raw_multinomial_p[i];
      }
    }
    if(PuzaDEF::Instance()->currentGeneratorType == GE_DIRICHLET){
      PuzaDEF::Instance()->GE_Dirichlet_K = dirichlet_k_Arg.getValue();
      PuzaDEF::Instance()->GE_Dirichlet_Alpha = dirichlet_alpha_Arg.getValue();
      vector<double> raw_dirichlet_base = dirichlet_base_Arg.getValue();
      if (PuzaDEF::Instance()->GE_Dirichlet_K == -1){
	cout << "In order to generate dirichlet random variables, we need K. (Assume scalar is 1.0)" << endl;
        exit(0);
      }
      if(raw_dirichlet_base.size()!=0){
	PuzaDEF::Instance()->GE_Dirichlet_Base = new double[raw_dirichlet_base.size()];
	for ( int i = 0; static_cast<unsigned int>(i) < raw_dirichlet_base.size(); i++ ){
	  PuzaDEF::Instance()->GE_Dirichlet_Base[i] = raw_dirichlet_base[i];
	}
      }
      else{
	PuzaDEF::Instance()->GE_Dirichlet_Base = new double[PuzaDEF::Instance()->GE_Dirichlet_K];
	for(int k=0; k < PuzaDEF::Instance()->GE_Dirichlet_K; k++){
	  PuzaDEF::Instance()->GE_Dirichlet_Base[k] = 1.0/ PuzaDEF::Instance()->GE_Dirichlet_K;
	}
      }
    }
  }  
}

#endif
